# 使用pymysql链接数据库
import pandas as pd
import pymysql


class SqlConnect:
    def __init__(self, host='localhost', user='root', password="123456", database='flask_trash_clasify'):
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.conn = None
        self.cursor = None

    def connect(self):
        self.conn = pymysql.connect(host=self.host, user=self.user, password=self.password, database=self.database)
        self.cursor = self.conn.cursor()

    def close(self):
        self.cursor.close()
        self.conn.close()

    def execute(self, sql, params=None):
        self.cursor.execute(sql, params)
        self.conn.commit()
        return True

    def fetchall(self):
        return self.cursor.fetchall()

    # 登录校验
    def check_login(self, username, password):
        self.connect()
        sql = "SELECT * FROM users WHERE username=%s AND password=%s"
        self.execute(sql, (username, password))
        result = self.fetchall()
        if len(result) == 1:
            return True
        else:
            return False

    # 注册用户
    def register_user(self, username, password):
        self.connect()
        # 检查用户名是否已存在
        sql = "SELECT * FROM users WHERE username=%s"
        self.execute(sql, (username,))
        result = self.fetchall()
        if len(result) == 1:
            return False
        else:
            # 插入用户信息
            sql = "INSERT INTO users(username, password) VALUES(%s, %s)"
            self.execute(sql, (username, password))

    # 读取所有用户数据为Dataframe
    def read_all_users(self):
        self.connect()
        sql = "SELECT * FROM users"
        # 执行使用pd.read_sql()方法
        df = pd.read_sql(sql, self.conn)
        return df

    # 根据username删除用户
    def delete_user(self, username):
        self.connect()
        sql = "DELETE FROM users WHERE username=%s"
        self.execute(sql, (username,))
        # 返回是否删除成功
        return True if self.cursor.rowcount == 1 else False

